import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn

def pool_nms(heat,kernel=3):
    '''非极大抑制'''
    pad=(kernel-1)//2

    hmax=nn.functional.max_pool2d(heat,(kernel,kernel),stride=1,padding=pad)#最大池化
    keep=(hmax==heat).float()
    return heat*keep

def decode_bbox(pred_hms,pred_whs,pred_offsets,image_size,threshold,cuda,topk=100):
    pred_hms=pool_nms(pred_hms)#对种类进行非极大抑制

    b, c, output_h,output_w=pred_hms.shape
    detects=[]
    for batch in range(b):
        heat_map=pred_hms[batch].permute(1,2,0).view([-1,c])
        pred_wh=pred_whs[batch].permute(1,2,0).view([-1,2])
        pred_offset=pred_offsets[batch].permute(1,2,0).view([-1,2])

        yv,xv=torch.meshgrid(torch.arange(0,output_h),torch.arange(0,output_w))

        xv,yv=xv.flatten().float(),yv.flatten().float()
        if cuda:
            xv=xv.cuda()
            yv=yv.cuda()

        class_conf,class_pred=torch.max(heat_map,dim=-1)
        mask=class_conf>threshold#最大置信度大于门限

        pred_wh_mask=pred_wh[mask]
        pred_offset_mask=pred_offset[mask]

        if len(pred_wh_mask)==0:
            detects.append([])
            continue

        xv_mask=torch.unsqueeze(xv[mask]+pred_offset_mask[...,0],-1)
        yv_mask = torch.unsqueeze(yv[mask] + pred_offset_mask[..., 0], -1)

        half_w, half_h = pred_wh_mask[..., 0:1] / 2, pred_wh_mask[..., 1:2] / 2
        bboxes = torch.cat([xv_mask - half_w, yv_mask - half_h, xv_mask + half_w, yv_mask + half_h], dim=1)
        bboxes[:, [0, 2]] /= output_w
        bboxes[:, [1, 3]] /= output_h
        detect = torch.cat(
            [bboxes, torch.unsqueeze(class_conf[mask], -1), torch.unsqueeze(class_pred[mask], -1).float()], dim=-1)

        arg_sort = torch.argsort(detect[:, -2], descending=True)
        detect = detect[arg_sort]

        detects.append(detect.cpu().numpy()[:topk])
    return detects
